# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import json
import logging
import random
import string
from typing import List, Union
from collections import namedtuple

import zlib

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.upbit_v1.kr_ws.ws_client import UpbitWsRequest
from coin.exchange.upbit_v1.kr_ws.upbit_parser import _strip_left

HashTimestamp = namedtuple('HashTimestamp', ['message_hash', 'timestamp'])

BOOK_UPDATE_INTERVAL_THRESHOLD = 20e9  # nanosecond


def _gen_message_hash(json_obj):
  data = json.dumps(json_obj)
  return zlib.crc32(data.encode())


def _get_native_symbol(json_obj):
  if 'code' in json_obj:
    native_symbol = json_obj['code']
  elif 'cd' in json_obj:
    native_symbol = json_obj['cd']
  else:
    raise ValueError(json_obj)
  return native_symbol


class UpbitSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)

    self._register_url(url=self.gen_address(), headers={'Origin': 'https://www.upbit.com'})
    self._publisher = None
    self._subscribed = False
    self._init_publisher(queue, group_id, worker_id)
    self._request = UpbitWsRequest(book_symbols=self._symbol, trade_symbols=self._symbol)
    self._request_str = self._request.as_request_str()
    self._last_hash_timestamp = {}

  def subscribe(self):
    assert not self._subscribed
    self.request()
    self._subscribed = True

  def request(self):
    self._ws.write_message(self._request_str)

  @staticmethod
  def gen_address():
    return "wss://crix-websocket.upbit.com/sockjs/%s/%s/websocket" % (
        "".join(random.choices(string.digits, k=3)),
        "".join(random.choices(string.ascii_lowercase, k=8)))

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('upbit_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('upbit_ws_trade_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'crixOrderbook': queue.get_writer(topic_id=topic_depth[1]),
        'crixTrade': queue.get_writer(topic_id=topic_trade[1]),
    }

    self._enable_topic_event_publisher(queue,
                                       'upbit_ws_topic_event', [topic_depth[0], topic_trade[0]])
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def _decode(self, payload_str):
    if payload_str.startswith("a"):
      # a?
      payload_str = _strip_left(payload_str, "a")
      payload = json.loads(payload_str)
      return payload
    elif payload_str.startswith("h"):
      self.request()
      # heartbeat
      return None
    elif payload_str.startswith("o"):
      self.subscribe()
      # open
      return None
    else:
      raise ValueError(payload_str)

  def on_message(self, timestamp: int, message: list):
    if self.stopped:
      return
    if message is None:
      return
    # not dict here... -_- it's a list...
    for payload_elem_str in message:
      message = json.loads(payload_elem_str)
      if 'type' in message:
        event_type = message['type']
      elif 'ty' in message:
        # New Upbit api use short style.
        event_type = message['ty']
      else:
        raise ValueError(message)
      if event_type == 'crixOrderbook':
        if not self.is_book_update_valid(timestamp, message):
          logging.error('Book update interval longer than %d sec, restart...' %
                        (BOOK_UPDATE_INTERVAL_THRESHOLD / 1e9))
          self.stop(trigger_close_callback=True)
          return
      try:
        self._publisher[event_type].write(timestamp, message)
      except KeyError:
        self._logger.error('Unknown event_type: %s' % event_type)

  # In Upbit, when feed is normal, usually book will be updated in every
  # second. For small coin, it could be every 5 - 10 seconds.
  # When feed is abnormal, book update is typically every 25 seconds. And
  # book value will be the same for that period and usually the value is
  # incorrect.
  # Here check the book content and update timestamp, if content hash is the
  # the same and it's updated after more than 20 seconds, consider shit happens
  # and restart the subscriber.
  def is_book_update_valid(self, timestamp, json_obj):
    native_symbol = _get_native_symbol(json_obj)
    last_hash_timestamp = self._last_hash_timestamp.get(native_symbol, None)

    if last_hash_timestamp is not None:
      message_hash = _gen_message_hash(json_obj)
      valid = ((message_hash != last_hash_timestamp.message_hash)
               or (timestamp - last_hash_timestamp.timestamp <= BOOK_UPDATE_INTERVAL_THRESHOLD))
      self._last_hash_timestamp[native_symbol] = HashTimestamp(message_hash, timestamp)
      return valid
    else:
      self._last_hash_timestamp[native_symbol] = HashTimestamp(_gen_message_hash(json_obj),
                                                               timestamp)
      return True
